import copy
import math

import numpy as np

from htm.Connection import Connection

RAT_BUMP_SIGMA = 0.18172


class LocationGuassian2D:

    def __init__(self, params):
        # Paramters
        self.cellsPerAxis = params['cellsPerAxis']
        self.scale = params['scale']
        self.bumpOverlapMethod = params['bumpOverlapMethod']
        self.baselineCellsPerAxis = params['baselineCellsPerAxis']
        self.orientation = params['orientation']
        # Initialization
        self.learn = True
        self.numOfCells = self.cellsPerAxis * self.cellsPerAxis
        self.A = np.linalg.inv(self.scale * np.array(
            [[math.cos(self.orientation), math.cos(self.orientation + np.radians(60.))],
             [math.sin(self.orientation), math.sin(self.orientation + np.radians(60.))]]))
        cellPhasesAxis = np.linspace(0., 1., self.cellsPerAxis, endpoint=False)
        self.cellPhases = np.array(
            [np.repeat(cellPhasesAxis, self.cellsPerAxis), np.tile(cellPhasesAxis, self.cellsPerAxis)])
        self.cellPhases += [[0.5 / self.cellsPerAxis], [0.5 / self.cellsPerAxis]]
        self.connection = Connection(params['connectionParams'])
        self.bumpSigma = RAT_BUMP_SIGMA * (self.baselineCellsPerAxis / float(self.cellsPerAxis))
        self.activeFiringRate = self.calFiringRate()
        # Mutable States
        self.bumpPhases = np.empty((2, 0), dtype="float")
        self.phaseDisplacement = np.empty((0, 2), dtype="float")
        self.cellsForActivePhases = np.empty(0, dtype="int")
        self.activeCells = np.empty(0, dtype="int")
        self.learningCells = np.empty(0, dtype="int")
        self.sensoryAssociatedCells = np.empty(0, dtype="int")
        self.activeSegments = np.empty(0, dtype="uint32")

    @staticmethod
    def gaussian(sig, d):
        return np.exp(-np.power(d, 2.) / (2 * np.power(sig, 2.)))

    def calFiringRate(self, minimumActiveDiameter=None):
        firingFieldDiameter = 2 * (1. / self.cellsPerAxis) * (2. / math.sqrt(3))
        if minimumActiveDiameter:
            firingFieldDiameter = max(firingFieldDiameter, minimumActiveDiameter)
        activeFiringRate = self.gaussian(self.bumpSigma, firingFieldDiameter / 2.)
        return activeFiringRate

    def reset(self):
        self.bumpPhases = np.empty((2, 0), dtype="float")
        self.phaseDisplacement = np.empty((0, 2), dtype="float")
        self.cellsForActivePhases = np.empty(0, dtype="int")
        self.activeCells = np.empty(0, dtype="int")
        self.learningCells = np.empty(0, dtype="int")
        self.sensoryAssociatedCells = np.empty(0, dtype="int")
        self.activeSegments = np.empty(0, dtype="uint32")

    def activateRandomLocation(self):
        rnd = np.array([np.random.random(2)])
        self.bumpPhases = rnd.T
        self.computeActiveCells()

    def computeActiveCells(self):
        cellExcitations = self.getCellExcitations()
        self.activeCells = np.where(cellExcitations >= self.activeFiringRate)[0]
        self.learningCells = np.where(cellExcitations == np.max(cellExcitations))[0]

    def getCellExcitations(self):
        cell_bump_positivePhaseDisplacement = np.mod(self.cellPhases.T[:, :, np.newaxis] - self.bumpPhases, 1.0)
        cell_direction_bump_phaseDisplacement = (cell_bump_positivePhaseDisplacement[:, np.newaxis, :, :] - np.array([[0, 0], [0, 1], [1, 0], [1, 1]])[:, :, np.newaxis])
        B = np.array([[np.cos(np.radians(0.)), np.cos(np.radians(60.))], [np.sin(np.radians(0.)), np.sin(np.radians(60.))]])
        cell_direction_bump_worldDisplacement = np.matmul(B, cell_direction_bump_phaseDisplacement)
        cell_direction_bump_distance = np.linalg.norm(cell_direction_bump_worldDisplacement, axis=-2)
        cell_bump_distance = np.amin(cell_direction_bump_distance, axis=1)
        cellExcitationsFromBumps = self.gaussian(self.bumpSigma, cell_bump_distance)
        if self.bumpOverlapMethod == "probabilistic":
            cellExcitations = 1. - np.prod(1. - cellExcitationsFromBumps, axis=1)
        elif self.bumpOverlapMethod == "sum":
            cellExcitations = np.sum(cellExcitationsFromBumps, axis=1)
        else:
            raise ValueError("Unrecognized bump overlap strategy", self.bumpOverlapMethod)
        return cellExcitations

    def sensoryCompute(self, anchorInput, anchorGrowthCandidates):
        if self.learn:
            self.sensoryComputeLearningMode(anchorGrowthCandidates)
        else:
            self.sensoryComputeInferenceMode(anchorInput)

    def sensoryComputeLearningMode(self, anchorInput):
        connectedOverlaps, potentialOverlaps = self.connection.computeActivity(anchorInput)
        activeSegments = self.connection.getActiveSegments(connectedOverlaps)
        matchingSegments = self.connection.getMatchingSegments(potentialOverlaps)
        cellsForActiveSegments = self.connection.mapSegmentsToCells(activeSegments)
        learningActiveSegments = activeSegments[np.in1d(cellsForActiveSegments, self.learningCells)]
        remainingCells = np.setdiff1d(self.learningCells, cellsForActiveSegments)
        candidateSegments = self.connection.filterSegmentsByCell(matchingSegments, remainingCells)
        cellsForCandidateSegments = np.unique(self.connection.mapSegmentsToCells(candidateSegments))
        learningMatchingSegments = self.connection.getBestSegmentPerCell(candidateSegments, potentialOverlaps)
        newSegmentCells = np.setdiff1d(remainingCells, cellsForCandidateSegments)
        self.connection.learnOnExisting(learningActiveSegments, anchorInput, anchorInput)
        self.connection.learnOnExisting(learningMatchingSegments, anchorInput, anchorInput)
        self.connection.learnOnNew(newSegmentCells, anchorInput)
        self.activeSegments = activeSegments
        self.sensoryAssociatedCells = self.learningCells

    def movementCompute(self, displacement, noiseFactor=0):
        if noiseFactor != 0:
            displacement = copy.deepcopy(displacement)
            xnoise = np.random.normal(0, noiseFactor)
            ynoise = np.random.normal(0, noiseFactor)
            displacement[0] += xnoise
            displacement[1] += ynoise
        phaseDisplacement = np.matmul(self.A, displacement)
        np.add(self.bumpPhases, phaseDisplacement[:, np.newaxis], out=self.bumpPhases)
        np.round(self.bumpPhases, decimals=9, out=self.bumpPhases)
        np.mod(self.bumpPhases, 1.0, out=self.bumpPhases)
        self.computeActiveCells()
        self.phaseDisplacement = phaseDisplacement

    def sensoryComputeInferenceMode(self, anchorInput):
        if len(anchorInput) == 0:
            return
        connectedOverlaps, _ = self.connection.computeActivity(anchorInput)
        activeSegments = self.connection.getActiveSegments(connectedOverlaps)
        sensorySupportedCells = np.unique(self.connection.mapSegmentsToCells(activeSegments))
        self.bumpPhases = self.cellPhases[:, sensorySupportedCells]
        self.computeActiveCells()
        self.activeSegments = activeSegments
        self.sensoryAssociatedCells = sensorySupportedCells

    def getNumOfCells(self):
        return self.numOfCells

    def getLearnableCells(self):
        return self.learningCells
